from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, List

from sqlglot import exp


@dataclass
class SplitQueues:
    ddl: List[exp.Expression]
    dml: List[exp.Expression]


def split_ddl_dml(expressions: Iterable[exp.Expression]) -> SplitQueues:
    ddl: List[exp.Expression] = []
    dml: List[exp.Expression] = []
    for e in expressions:
        # Include generic command nodes (e.g., USE/SET) in DDL bucket
        cls_name = getattr(e, '__class__', type('x', (), {})).__name__.upper()
        is_generic_cmd = isinstance(e, getattr(exp, 'Command', tuple())) or cls_name in ("USE", "SET", "COMMAND")
        if isinstance(e, (exp.Create, exp.Drop, exp.Alter, exp.Comment)) or is_generic_cmd:  # DDL bucket
            ddl.append(e)
        else:
            dml.append(e)
    return SplitQueues(ddl=ddl, dml=dml)


